{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agents\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup LLM and Embedding Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4\", temperature=0.1)\n",
    "embed_model = OpenAIEmbedding()\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents and Tools usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.core.agent.workflow import (\n",
    "    FunctionAgent,\n",
    "    ReActAgent,\n",
    ")\n",
    "\n",
    "from IPython.display import display, HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two integers and returns the result integer\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "add_tool = FunctionTool.from_defaults(fn=add)\n",
    "subtract_tool = FunctionTool.from_defaults(fn=subtract)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With ReAct Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = ReActAgent(\n",
    "    tools=[multiply_tool, add_tool, subtract_tool],\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"What is (26 * 2) + 2024?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">2076</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With Function Calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = FunctionAgent(\n",
    "    tools=[multiply_tool, add_tool, subtract_tool],\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"What is (26 * 2) + 2024?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">assistant: The result of (26 * 2) + 2024 is 2076.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(f'<p style=\"font-size:20px\">{response}</p>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent with RAG Query Engine Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download Data\n",
    "\n",
    "We will use `Uber-2021` and `Lyft-2021` 10K SEC filings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O './uber_2021.pdf'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/lyft_2021.pdf' -O './lyft_2021.pdf'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "uber_docs = SimpleDirectoryReader(input_files=[\"./uber_2021.pdf\"]).load_data()\n",
    "lyft_docs = SimpleDirectoryReader(input_files=[\"./lyft_2021.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build RAG on uber and lyft docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "uber_index = VectorStoreIndex.from_documents(uber_docs)\n",
    "uber_query_engine = uber_index.as_query_engine(similarity_top_k=3)\n",
    "\n",
    "lyft_index = VectorStoreIndex.from_documents(lyft_docs)\n",
    "lyft_query_engine = lyft_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = uber_query_engine.query(\"What are the investments of Uber in 2021?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">In 2021, Uber made significant investments to expand its international operations and compete with local and other global competitors. This included the acquisitions of Careem and Cornershop. Additionally, Uber continued to develop new technologies to enhance existing offerings and services, and to expand the range of its offerings through research and development. They also launched Uber One in the United States, a cross-platform membership program that brings together the best of Uber.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = lyft_query_engine.query(\"What are lyft investments in 2021?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">In 2021, Lyft continued to invest in expanding its network of Light Vehicles and Lyft Autonomous, focusing on the deployment and scaling of third-party self-driving technology on the Lyft network. The company also made a commitment to reach 100% electric vehicles on the Lyft network by the end of 2030. Additionally, Lyft completed a transaction with Woven Planet, a subsidiary of Toyota Motor Corporation, for the divestiture of certain assets related to its self-driving vehicle division, Level 5.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(f'<p style=\"font-size:20px\">{response.response}</p>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `FunctionAgent` with RAG QueryEngineTools.\n",
    "\n",
    "Here we use `Fuction Calling` capabilities of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "query_engine_tools = [\n",
    "    QueryEngineTool(\n",
    "        query_engine=lyft_query_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"lyft_10k\",\n",
    "            description=\"Provides information about Lyft financials for year 2021\",\n",
    "        ),\n",
    "    ),\n",
    "    QueryEngineTool(\n",
    "        query_engine=uber_query_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"uber_10k\",\n",
    "            description=\"Provides information about Uber financials for year 2021\",\n",
    "        ),\n",
    "    ),\n",
    "]\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=query_engine_tools,\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"What are the investments of Uber in 2021?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">assistant: In 2021, Uber's investments primarily consisted of money market funds, cash deposits, U.S. government and agency securities, and investment-grade corporate debt securities. Their investment policy aims to preserve capital and meet liquidity requirements without significantly increasing risk. As of December 31, 2021, Uber had cash and cash equivalents including restricted cash and cash equivalents totaling $7.8 billion. They also hold investments in other companies, including minority-owned, privately-held affiliates and recently public companies. The carrying value of these investments was $12.6 billion as of December 31, 2021.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(f'<p style=\"font-size:20px\">{response}</p>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await agent.run(\"What are lyft investments in 2021?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p style=\"font-size:20px\">assistant: In 2021, Lyft's investments included cash and cash equivalents, short-term investments, and restricted investments. Cash and cash equivalents included certificates of deposits, commercial paper, and corporate bonds that have an original maturity of 90 days or less and are readily convertible to known amounts of cash. Short-term investments were comprised of commercial paper, certificates of deposit, and corporate bonds, which mature in twelve months or less. \n",
       "\n",
       "Restricted investments were comprised of debt security investments in commercial paper, certificates of deposit, corporate bonds, and U.S. government securities which are held in trust accounts at third-party financial institutions pursuant to certain contracts with insurance providers. \n",
       "\n",
       "The company also had investments in non-marketable equity securities, which are measured at cost, with remeasurements to fair value only upon the occurrence of observable transactions for identical or similar investments of the same issuer or impairment.</p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(f'<p style=\"font-size:20px\">{response}</p>'))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "mistralai22b",
   "language": "python",
   "name": "mistralai22b"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
